package com.example.demo.controller;

import cn.hutool.cache.Cache;
import cn.hutool.cache.CacheUtil;
import cn.hutool.core.lang.UUID;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.example.demo.app.LoveApp;
import com.example.demo.model.dto.ai.LoveReport;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;


@RestController
@RequestMapping("/helloworld")
public class HelloworldController {

    @Resource
    LoveApp loveApp;
    private static final String DEFAULT_PROMPT = "你是一个博学的智能聊天助手，请根据用户提问回答！";

    // 用一个缓存来存储每个消息的事件列表
// Key: messageId, Value: List of ServerSentEvent
    private final Cache<String, List<ServerSentEvent<String>>> messageCache = CacheUtil.newLRUCache(100);


    private final ChatClient dashScopeChatClient;

    public HelloworldController(ChatClient.Builder chatClientBuilder) {
        // 初始化基于内存的对话记忆
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(new InMemoryChatMemoryRepository())
                .maxMessages(20)
                .build();
        this.dashScopeChatClient = chatClientBuilder
                .defaultSystem(DEFAULT_PROMPT)
                // 实现 Chat Memory 的 Advisor
                // 在使用 Chat Memory 时，需要指定对话 ID，以便 Spring AI 处理上下文。
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build()
                )
                // 实现 Logger 的 Advisor
                .defaultAdvisors(
                        new SimpleLoggerAdvisor()
                )
                // 设置 ChatClient 中 ChatModel 的 Options 参数
                .defaultOptions(
                        DashScopeChatOptions.builder()
                                .withTopP(0.7)
                                .build()
                )
                .build();
    }

    @GetMapping("/simple/chat")
    public String simpleChat(String query) {
        return dashScopeChatClient.prompt(query).call().content();
    }

    @GetMapping("/simple1/chat")
    public LoveReport simpleChat1(String query) {
        return loveApp.doChat(query, "1");
    }


    /**
     * ChatClient 流式调用
     */
    @GetMapping("/stream/chat")
    public Flux<String> streamChat(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？") String query, HttpServletResponse response) {

        response.setCharacterEncoding("UTF-8");
        return dashScopeChatClient.prompt(query).stream().content();
    }

    /**
     * ChatClient 流式调用，支持断点续传。
     * <p>
     * 1.  **首次请求**:
     * - GET /helloworld/stream/chat/simple-resume?query=你好
     * - 服务器返回 "start" 事件 (data 为 messageId)，然后直接开始返回 "content" 事件。
     * <p>
     * 2.  **重连**:
     * - 使用首次请求收到的 messageId 和最后一个事件的 ID (Last-Event-ID) 来请求。
     * - GET /helloworld/stream/chat/simple-resume?query=你好&messageId={messageId}
     * - Header: Last-Event-ID: {lastEventId}
     * - 服务器将从缓存中续传丢失的事件。
     *
     * @param query       用户查询
     * @param messageId   会话消息ID，用于重连
     * @param lastEventId 客户端收到的最后一个事件ID，用于重连
     * @param response    HttpServletResponse
     * @return ServerSentEvent 的响应流
     */
    @GetMapping("/stream/chat/simple-resume")
    public Flux<ServerSentEvent<String>> simpleStreamChatWithResume(
            @RequestParam(value = "query") String query,
            @RequestParam(value = "messageId", required = false) String messageId,
            @RequestHeader(name = "Last-Event-ID", required = false) String lastEventId,
            HttpServletResponse response) {

        response.setCharacterEncoding("UTF-8");
        response.setContentType("text/event-stream");

        // --- 断点续传逻辑 (与之前的版本完全相同) ---
        if (messageId != null && !messageId.isEmpty() && lastEventId != null && !lastEventId.isEmpty()) {
            List<ServerSentEvent<String>> cachedEvents = messageCache.get(messageId);
            if (cachedEvents != null) {
                try {
                    int lastId = Integer.parseInt(lastEventId);
                    List<ServerSentEvent<String>> eventsToResend = cachedEvents.stream()
                            .filter(event -> {
                                try {
                                    return Integer.parseInt(event.id()) > lastId;
                                } catch (NumberFormatException e) {
                                    return false;
                                }
                            })
                            .collect(Collectors.toList());

                    if (!eventsToResend.isEmpty()) {
                        return Flux.fromIterable(eventsToResend);
                    }
                } catch (NumberFormatException e) {
                    // lastEventId 格式错误，当作新请求处理
                }
            }
        }
        // --- 新的流式请求 ---
        final String currentMessageId = (messageId == null || messageId.isEmpty()) ? UUID.randomUUID().toString() : messageId;
        final List<ServerSentEvent<String>> events = new CopyOnWriteArrayList<>();
        messageCache.put(currentMessageId, events);

        final AtomicInteger eventIdCounter = new AtomicInteger(0);

        // 1. "start" 事件：通知客户端流开始，并发送 messageId
        ServerSentEvent<String> startEvent = ServerSentEvent.<String>builder()
                .id(String.valueOf(eventIdCounter.getAndIncrement()))
                .event("start")
                .data(currentMessageId)
                .build();

        // 将 start 事件加入缓存
        events.add(startEvent);

        // 2. "content" 事件流
        Flux<ServerSentEvent<String>> contentStream = dashScopeChatClient.prompt(query).stream().content()
                .map(chunk -> ServerSentEvent.<String>builder()
                        .id(String.valueOf(eventIdCounter.getAndIncrement()))
                        .event("content")
                        .data(chunk)
                        .build())
                .doOnNext(events::add); // 缓存每个 content 事件

        // 3. "done" 事件
        Flux<ServerSentEvent<String>> doneStream = Flux.create(sink -> {
            ServerSentEvent<String> doneEvent = ServerSentEvent.<String>builder()
                    .id(String.valueOf(eventIdCounter.getAndIncrement()))
                    .event("done")
                    .data("Stream completed.")
                    .build();
            events.add(doneEvent); // 缓存 done 事件
            sink.next(doneEvent);
            sink.complete();
        });

        // 合并事件流: start -> content -> done
        return Flux.just(startEvent)
                .concatWith(contentStream)
                .concatWith(doneStream);
    }

    /**
     * ChatClient 使用自定义的 Advisor 实现功能增强.
     * eg:
     * http://127.0.0.1:18080/helloworld/advisor/chat/123?query=你好，我叫牧生，之后的会话中都带上我的名字
     * 你好，牧生！很高兴认识你。在接下来的对话中，我会记得带上你的名字。有什么想聊的吗？
     * http://127.0.0.1:18080/helloworld/advisor/chat/123?query=我叫什么名字？
     * 你叫牧生呀。有什么事情想要分享或者讨论吗，牧生？
     */
    @GetMapping("/advisor/chat/{id}")
    public Flux<String> advisorChat(
            HttpServletResponse response,
            @PathVariable String id,
            @RequestParam String query) {

        response.setCharacterEncoding("UTF-8");

        return this.dashScopeChatClient.prompt(query)
                .advisors(
                        // TODO
//						a -> a
//								.param(CHAT_MEMORY_CONVERSATION_ID_KEY, id)
//								.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100)
                ).stream().content();
    }


}